import os
import folder_paths
import logging  # 添加logging导入

# 配置日志
logger = logging.getLogger('PSDModifier')

# 定义新的数据类型用于传递文件路径列表
if not hasattr(folder_paths, 'IMAGE_PATH_LIST'):
    folder_paths.IMAGE_PATH_LIST = "IMAGE_PATH_LIST"

class ListImagesInDirectory:
    """
    从目录列出图像文件节点
    扫描指定文件夹内的所有图片文件，提取每张图片的完整路径及名称信息，
    并将文件路径列表向下游节点传输
    """
    # 添加静态缓存字典，用于存储已扫描的目录结果
    _cache = {}
    # 定义节点的输入类型
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "directory_path": ("STRING", {"default": "", "multiline": False}),  # 文件夹路径输入
            },
            "optional": {
                "recursive": ("BOOLEAN", {"default": False}),  # 是否递归子文件夹
                "filter_extensions": ("STRING", {"default": "png,jpg,jpeg,bmp,gif", "multiline": False}),  # 过滤的文件扩展名
                "max_images": ("INT", {"default": 100, "min": 1, "max": 1000, "step": 1}),  # 最大加载的图像数量
                "reset_cache": ("BOOLEAN", {"default": False}),  # 重置缓存参数
            }
        }

    # 定义节点的输出类型
    RETURN_TYPES = (folder_paths.IMAGE_PATH_LIST, "INT", "STRING")
    RETURN_NAMES = ("IMAGE_PATHS", "image_count", "debug_info")
    FUNCTION = "list_images"
    CATEGORY = "image/utils"
    OUTPUT_IS_LIST = (False, False, False)  # 所有输出都是单个值

    def list_images(self, directory_path="", recursive=False, filter_extensions="png,jpg,jpeg,bmp,gif", max_images=100, reset_cache=False):
        """
        从指定目录扫描并列出图像文件路径，支持缓存机制和缓存重置功能
        
        参数:
            directory_path: 要扫描的文件夹路径
            recursive: 是否递归子文件夹
            filter_extensions: 要过滤的文件扩展名，逗号分隔
            max_images: 最大获取的图像文件数量
            reset_cache: 是否重置当前目录的缓存
        
        返回:
            tuple: 包含图像路径列表、图像总数和调试信息的元组
        """
        # 准备返回值
        image_paths = []
        image_count = 0
        debug_info = ""
        
        # 检查目录是否存在
        if not directory_path or not os.path.exists(directory_path):
            logger.error(f"ListImagesInDirectory: 目录不存在: {directory_path}")
            debug_info = f"错误: 目录不存在: {directory_path}"
            return (image_paths, image_count, debug_info)
            
        if not os.path.isdir(directory_path):
            logger.error(f"ListImagesInDirectory: 不是有效的目录: {directory_path}")
            debug_info = f"错误: 不是有效的目录: {directory_path}"
            return (image_paths, image_count, debug_info)
        
        # 生成缓存键
        cache_key = (directory_path, recursive, filter_extensions, max_images)
        
        # 检查是否需要重置缓存
        if reset_cache and cache_key in ListImagesInDirectory._cache:
            del ListImagesInDirectory._cache[cache_key]
            logger.info(f"ListImagesInDirectory: 已重置目录 {directory_path} 的缓存")
        
        # 检查缓存中是否已有结果
        if cache_key in ListImagesInDirectory._cache:
            image_paths, image_count, debug_info = ListImagesInDirectory._cache[cache_key]
            logger.info(f"ListImagesInDirectory: 从缓存获取目录 {directory_path} 的结果")
            return (image_paths, image_count, debug_info)
            
        # 解析文件扩展名过滤器
        valid_extensions = [ext.strip().lower() for ext in filter_extensions.split(',') if ext.strip()]
        
        # 准备搜索函数
        def is_valid_image_file(file_path):
            """检查文件是否为有效的图像文件"""
            if not os.path.isfile(file_path):
                return False
            
            # 检查文件扩展名
            _, ext = os.path.splitext(file_path)
            ext = ext.lstrip('.').lower()
            if valid_extensions and ext not in valid_extensions:
                return False
            
            # 尝试打开文件以确认它是图像
            try:
                from PIL import Image
                with Image.open(file_path) as img:
                    # 验证是否为有效图像
                    img.verify()
                return True
            except Exception as e:
                logger.warning(f"ListImagesInDirectory: 不是有效的图像文件 {file_path}: {str(e)}")
                return False
                
        # 搜索图像文件
        image_files = []
        
        if recursive:
            # 递归搜索子文件夹
            for root, dirs, files in os.walk(directory_path):
                for file in files:
                    file_path = os.path.join(root, file)
                    if is_valid_image_file(file_path):
                        image_files.append(file_path)
                        if len(image_files) >= max_images:
                            break
                if len(image_files) >= max_images:
                    break
        else:
            # 仅在当前目录搜索
            try:
                for file in os.listdir(directory_path):
                    file_path = os.path.join(directory_path, file)
                    if is_valid_image_file(file_path):
                        image_files.append(file_path)
                        if len(image_files) >= max_images:
                            break
            except Exception as e:
                logger.error(f"ListImagesInDirectory: 读取目录内容失败: {str(e)}")
                debug_info = f"错误: 读取目录内容失败: {str(e)}"
                return (image_paths, image_count, debug_info)
        
        # 处理搜索结果
        image_count = len(image_files)
        image_paths = image_files  # 将找到的图像文件路径列表赋值给输出
        
        # 更新缓存
        debug_info = f"成功: 找到 {image_count} 张图像文件"
        ListImagesInDirectory._cache[cache_key] = (image_paths, image_count, debug_info)
        logger.info(f"ListImagesInDirectory: {debug_info}")
        
        return (image_paths, image_count, debug_info)

# 注册节点
NODE_CLASS_MAPPINGS = {
    "ListImagesInDirectory": ListImagesInDirectory
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "ListImagesInDirectory": "List Images In Directory"
}